{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import what.utils.logger as log\n",
    "\n",
    "logger = log.get_logger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_ssd = \"models/mobilenet-v1-ssd-mp-0_675.pth\"\n",
    "checkpoint_folder = \"models/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Training Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "num_workers = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fiftyone as fo\n",
    "import fiftyone.zoo as foz\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from what.models.detection.datasets.fiftyone import FiftyOneDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from what.models.detection.ssd.ssd.preprocessing import TrainAugmentation, TestTransform\n",
    "from what.models.detection.ssd.ssd import mobilenetv1_ssd_config\n",
    "from what.models.detection.ssd.ssd.ssd import MatchPrior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load MobileNetSSD configuration\n",
    "config              = mobilenetv1_ssd_config\n",
    "train_transform     = TrainAugmentation(config.image_size, config.image_mean, config.image_std)\n",
    "target_transform    = MatchPrior(config.priors, config.center_variance, config.size_variance, 0.5)\n",
    "test_transform      = TestTransform(config.image_size, config.image_mean, config.image_std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split 'train' already downloaded\n",
      "Ignoring unsupported parameter 'label_types' for importer type <class 'fiftyone.utils.data.importers.FiftyOneImageDetectionDatasetImporter'>\n",
      "Loading existing dataset 'voc-2012-train'. To reload from disk, either delete the existing dataset or provide a custom `dataset_name` to use\n"
     ]
    }
   ],
   "source": [
    "# Load the data from fiftyone\n",
    "voc_2012_train = foz.load_zoo_dataset(\"voc-2012\", split=\"train\", label_types=[\"detections\"],)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "session = fo.launch_app(voc_2012_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO] 2022-01-20 15:15:28,574 __main__: Train dataset size: 5717\n"
     ]
    }
   ],
   "source": [
    "# Convert to pytorch dataloader\n",
    "train_dataset = FiftyOneDataset(voc_2012_train, \n",
    "                                foz.load_zoo_dataset_info(\"voc-2012\").classes, \n",
    "                                transform=train_transform,\n",
    "                                target_transform=target_transform)\n",
    "\n",
    "# Load Training Dataset from Local Disk\n",
    "# train_dataset = VOCDataset(\"examples/VOC2012\", transform=train_transform,\n",
    "#                              target_transform=target_transform)\n",
    "\n",
    "train_loader  = DataLoader(train_dataset, batch_size,\n",
    "                          num_workers=num_workers,\n",
    "                          shuffle=True)\n",
    "\n",
    "logger.info(\"Train dataset size: {}\".format(len(train_dataset)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Validation Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split 'train' already downloaded\n",
      "Loading existing dataset 'voc-2007-train'. To reload from disk, either delete the existing dataset or provide a custom `dataset_name` to use\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO] 2022-01-20 15:15:30,608 __main__: validation dataset size: 2501\n"
     ]
    }
   ],
   "source": [
    "# Load Validation Dataset from FiftyOne (use voc-2007 train as validation here)\n",
    "val_dataset  = FiftyOneDataset(foz.load_zoo_dataset(\"voc-2007\", split=\"train\"), \n",
    "                                        foz.load_zoo_dataset_info(\"voc-2007\").classes,\n",
    "                                        transform=test_transform,\n",
    "                                        target_transform=target_transform)\n",
    "\n",
    "# Load Validation Dataset from Local Disk\n",
    "# val_dataset = VOCDataset(\"VOC2007/\", transform=test_transform,\n",
    "#                              target_transform=target_transform, is_test=True)\n",
    "\n",
    "val_loader = DataLoader(val_dataset, batch_size,\n",
    "                        num_workers=num_workers,\n",
    "                        shuffle=False)\n",
    "\n",
    "logger.info(\"validation dataset size: {}\".format(len(val_dataset)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Mobilenet V1 SSD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from what.models.detection.ssd.mobilenet_v1_ssd import MobileNetV1SSD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO] 2022-01-20 15:15:32,661 what.models.detection.ssd.mobilenet_v1_ssd: Init from pretrained ssd models/mobilenet-v1-ssd-mp-0_675.pth\n",
      "[INFO] 2022-01-20 15:15:34,147 what.models.detection.ssd.mobilenet_v1_ssd: Start training using CosineAnnealingLR scheduler.\n",
      "/home/wuhanstudio/whitebox-adversarial-toolbox/what/models/detection/ssd/transforms/transforms.py:243: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  mode = random.choice(self.sample_options)\n",
      "[INFO] 2022-01-20 15:15:41,215 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 10, Average Loss: 19.7620, Average Regression Loss 3.9450, Average Classification Loss: 15.8170\n",
      "[INFO] 2022-01-20 15:15:46,461 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 20, Average Loss: 14.0563, Average Regression Loss 2.9937, Average Classification Loss: 11.0626\n",
      "[INFO] 2022-01-20 15:15:51,890 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 30, Average Loss: 11.6577, Average Regression Loss 2.9390, Average Classification Loss: 8.7187\n",
      "[INFO] 2022-01-20 15:15:57,324 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 40, Average Loss: 9.7512, Average Regression Loss 2.7660, Average Classification Loss: 6.9852\n",
      "[INFO] 2022-01-20 15:16:02,809 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 50, Average Loss: 8.6162, Average Regression Loss 2.6561, Average Classification Loss: 5.9601\n",
      "[INFO] 2022-01-20 15:16:08,262 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 60, Average Loss: 7.9420, Average Regression Loss 2.4775, Average Classification Loss: 5.4645\n",
      "[INFO] 2022-01-20 15:16:13,942 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 70, Average Loss: 7.8411, Average Regression Loss 2.6946, Average Classification Loss: 5.1465\n",
      "[INFO] 2022-01-20 15:16:19,312 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 80, Average Loss: 7.2090, Average Regression Loss 2.2292, Average Classification Loss: 4.9799\n",
      "[INFO] 2022-01-20 15:16:24,890 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 90, Average Loss: 7.0581, Average Regression Loss 2.2933, Average Classification Loss: 4.7648\n",
      "[INFO] 2022-01-20 15:16:30,210 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 100, Average Loss: 6.8130, Average Regression Loss 2.1301, Average Classification Loss: 4.6828\n",
      "[INFO] 2022-01-20 15:16:35,586 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 110, Average Loss: 6.5294, Average Regression Loss 2.0344, Average Classification Loss: 4.4950\n",
      "[INFO] 2022-01-20 15:16:41,162 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 120, Average Loss: 6.5411, Average Regression Loss 1.9626, Average Classification Loss: 4.5785\n",
      "[INFO] 2022-01-20 15:16:46,779 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 130, Average Loss: 6.2685, Average Regression Loss 1.8621, Average Classification Loss: 4.4064\n",
      "[INFO] 2022-01-20 15:16:52,010 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 140, Average Loss: 6.3041, Average Regression Loss 1.9504, Average Classification Loss: 4.3536\n",
      "[INFO] 2022-01-20 15:16:57,610 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 150, Average Loss: 6.1556, Average Regression Loss 1.7948, Average Classification Loss: 4.3608\n",
      "[INFO] 2022-01-20 15:17:03,056 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 160, Average Loss: 6.2693, Average Regression Loss 1.9008, Average Classification Loss: 4.3685\n",
      "[INFO] 2022-01-20 15:17:08,486 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Step: 170, Average Loss: 6.2064, Average Regression Loss 1.8248, Average Classification Loss: 4.3815\n",
      "[INFO] 2022-01-20 15:17:32,680 what.models.detection.ssd.mobilenet_v1_ssd: Epoch: 0, Validation Loss: 5.7942, Validation Regression Loss 1.6806, Validation Classification Loss: 4.1136\n",
      "[INFO] 2022-01-20 15:17:32,732 what.models.detection.ssd.mobilenet_v1_ssd: Saved model models/mobilenet-v1-ssd-Epoch-0-Loss-5.794220924377441.pth\n"
     ]
    }
   ],
   "source": [
    "# Create SSD network and load pretrained base net.\n",
    "model = MobileNetV1SSD(is_test=False, class_names=train_dataset.classes)\n",
    "\n",
    "model.train(train_loader, val_loader, device=device, num_epochs=1, debug_steps=10, validation_epochs=1,\n",
    "           pretrained_ssd = pretrained_ssd, checkpoint_folder = checkpoint_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
